from fastapi import APIRouter, UploadFile, File, HTTPException, WebSocket
from uuid import uuid4
import json
import time
import os
import shutil
from schemas import UploadResponse, StartJobRequest, JobResponse, JobStatus, StopJobResponse
from tasks.training_tasks import run_task
from core.redis_client import redis_client

router = APIRouter(prefix="/jobs", tags=["jobs"])

@router.post("/upload", response_model=UploadResponse)
async def upload_file(file: UploadFile = File(...)):
    upload_dir = os.path.join(os.getcwd(), "uploadFiles")
    os.makedirs(upload_dir, exist_ok=True)
    file_path = os.path.join(upload_dir, file.filename)
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    return UploadResponse(success=True, filePath=file_path, fileName=file.filename)

@router.post("/start", response_model=JobResponse)
async def start_job(request: StartJobRequest):
    job_id = str(uuid4())
    # 初始化Redis状态
    redis_client.hset(f"job:{job_id}", mapping={
        "jobId": job_id,
        "taskType": request.taskType,
        "status": "queued",
        "createdAt": str(int(time.time())),
        "stopRequested": "0"
    })
    # 启动Celery任务
    run_task.delay(job_id, request.taskType, request.dataFile, request.params, request.modelId)
    return JobResponse(success=True, jobId=job_id)

@router.get("/{job_id}/status", response_model=JobStatus)
async def get_job_status(job_id: str):
    data = redis_client.hgetall(f"job:{job_id}")
    if not data:
        raise HTTPException(status_code=404, detail="Job not found")
    progress = json.loads(data.get("progress", "{}")) if data.get("progress") else None
    results = json.loads(data.get("result", "[]")) if data.get("result") else None
    message = data.get("message", "任务进行中")
    return JobStatus(
        jobId=job_id,
        status=data["status"],
        taskType=data["taskType"],
        progress=progress,
        results=results,
        message=message
    )

@router.post("/{job_id}/stop", response_model=StopJobResponse)
async def stop_job(job_id: str):
    if not redis_client.exists(f"job:{job_id}"):
        raise HTTPException(status_code=404, detail="Job not found")
    redis_client.hset(f"job:{job_id}", mapping={
        "stopRequested": "1",
        "status": "stopping"
    })
    redis_client.publish(f"job-updates:{job_id}", json.dumps({"status": "stopping"}))
    return StopJobResponse(success=True, message=f"任务 {job_id} 已中止。")

# WebSocket for 实时推送
@router.websocket("/{job_id}/ws")
async def websocket_endpoint(websocket: WebSocket, job_id: str):
    await websocket.accept()
    pubsub = redis_client.pubsub()
    pubsub.subscribe(f"job-updates:{job_id}")
    try:
        for message in pubsub.listen():
            if message.get("type") == "message":
                await websocket.send_text(message["data"])
    finally:
        pubsub.unsubscribe()
        await websocket.close()


        